import moviepy.editor as mpy
import matplotlib.pyplot as plt, numpy as np, seaborn as sns, tqdm, pathlib, pims, cv2 as cv, xarray as xf, pandas as pd, scipy.ndimage as ndimage

vid = pims.Video('observations/4_data_settling_velocities/videos/15-44-35_20-Jun-19_settling_velocity.mov')
# img = vid[1440][:,150:-150]
# cv.imwrite('poster_pic2.png', img)
nframes = vid._len
# nframes = 3000
start = 0
cent = np.zeros((2,nframes))
dt = 1/120
pix2mm = 200 / vid[0][:,120:-20].shape[1]

for i in tqdm.tqdm(range(start,start+nframes)):
    img = vid[i]
    img = img[:,600:850]
    img = cv.cvtColor(img, cv.COLOR_RGB2GRAY)
    img = cv.medianBlur(img, 7)
    img = cv.adaptiveThreshold(img, 255, cv.ADAPTIVE_THRESH_GAUSSIAN_C, cv.THRESH_BINARY, 201, -41)
    img = cv.medianBlur(img, 5)
    kernel = np.ones((3,3),np.uint8)
    img = cv.morphologyEx(img, cv.MORPH_OPEN, kernel, iterations=2)
    img = cv.morphologyEx(img, cv.MORPH_CLOSE, kernel, iterations=4)
    img[img>0] = 1
    if img.max() > 0:
        cent[:,i-start] = ndimage.measurements.center_of_mass(img)
    else:
        cent[:,i-start] = np.nan

y_pos = (-cent[1,:]) * pix2mm
x_pos = cent[0,:] * pix2mm
x_vel = np.zeros_like(x_pos)
y_vel = np.zeros_like(y_pos)
x_vel[0] = (x_pos[1] - x_pos[0]) / dt
x_vel[1] = (x_pos[2] - x_pos[1]) / dt
x_vel[-2] = (x_pos[-2] - x_pos[-3]) / dt
x_vel[-1] = (x_pos[-1] - x_pos[-2]) / dt
y_vel[0] = (y_pos[1] - y_pos[0]) / dt
y_vel[1] = (y_pos[2] - y_pos[1]) / dt
y_vel[-2] = (y_pos[-2] - y_pos[-3]) / dt
y_vel[-1] = (y_pos[-1] - y_pos[-2]) / dt
for ii in tqdm.tqdm(range(2,nframes-2)):
    x_vel[ii] = (-x_pos[ii+2] + 8*x_pos[ii+1] - 8*x_pos[ii-1] + x_pos[ii-2]) / (12*dt)
    y_vel[ii] = (-y_pos[ii+2] + 8*y_pos[ii+1] - 8*y_pos[ii-1] + y_pos[ii-2]) / (12*dt)

# vel_mag = np.sqrt(x_vel**2 + y_vel**2)
vel_mag = np.sqrt(y_vel**2)
time = (np.ones_like(vel_mag) * dt).cumsum()
df = pd.DataFrame({'time': time, 'v': vel_mag})
df['v_smooth'] = df.v.interpolate(limit=11).rolling(window=31, min_periods=11, center=True).mean()
df['particleId'] = np.zeros_like(df.v.values)
ii = 0
particleId = 0
start_list = np.where(np.isfinite(df.v_smooth.values))[0]
while ii < df.shape[0]:
    if np.isfinite(df.v_smooth.values[ii]):
        df['particleId'][ii] = particleId
        ii += 1
    else:
        particleId += 1
        ilist = start_list[start_list > ii]
        if len(ilist) > 0:
            ii = ilist[0]
        else:
            break

count = 0
plt.figure(count)
dfs = []
for name, group in df.groupby('particleId'):
    if (group.v.std() < 40) & (group.v.mean() > 100):
        count += 1
        plt.plot(group.v.values, '-')
        # plt.title(group.v.mean())
        dfs.append(group)
plt.ylim(0, 1000)
v_mean = pd.concat(dfs).groupby('particleId').mean().v
v_std = pd.concat(dfs).groupby('particleId').std().v
plt.title('mean v: %g +/- %g' % (v_mean.mean(), v_std.mean()))
plt.savefig('observations/4_data_settling_velocities/individual_grain_material_processing_codes/terminal_settling_tracks_gb.pdf')
df = pd.DataFrame(dict(vel = v_mean, vel_e = v_std))
df.to_csv('observations/4_data_settling_velocities/gb_settling_velocities.csv')


# fps = 60
# num_images = nframes
# # num_images = 3000
# duration = num_images/fps
#
# def make_frame(t):
#     img = vid[int(t*fps)]
#     img = img[:,500:850]
#     img = cv.cvtColor(img, cv.COLOR_RGB2GRAY)
#     img = cv.medianBlur(img, 7)
#     img = cv.adaptiveThreshold(img, 255, cv.ADAPTIVE_THRESH_GAUSSIAN_C, cv.THRESH_BINARY, 201, -41)
#     img = cv.medianBlur(img, 5)
#     kernel = np.ones((3,3),np.uint8)
#     img = cv.morphologyEx(img, cv.MORPH_OPEN, kernel, iterations=2)
#     img = cv.morphologyEx(img, cv.MORPH_CLOSE, kernel, iterations=4)
#     img = cv.cvtColor(img, cv.COLOR_GRAY2RGB)
#     return img
#
# animation = mpy.VideoClip(make_frame, duration=duration)
# animation.write_videofile('observations/4_data_settling_velocities/videos/15-44-35_20-Jun-19_settling_velocity_converted.mp4', fps=fps) # export as video
